#ifndef AST_HPP
#define AST_HPP

#include <vector>
#include <string>
using namespace std;

class Token;
class Visitor;
class SymbolTable;
class TypeTable;
class Type;
class Dumper;

class Context {
public:
    static const int Load = 0;
    static const int Store = 1;
};

class Node {
friend Dumper;
private:
    Node* _parent;
    int   _silent;

protected:
    Type* _type;

public:
    Node() : _parent(nullptr), _type(nullptr), _silent(0)   {}
    virtual Node* accept(Visitor* v) = 0;

    void set_parent(Node* p)    { _parent = p; }
    Node* parent()              { return _parent; }

    int silent() { return _silent; }
    void set_silent(int s) { _silent = s; }

    virtual Type* get_type(TypeTable* type_table);
    void set_type(Type* t) { _type = t; }

    virtual string to_string() { return "Node"; }
};

enum OpType {
    AST_OP_ADD,
    AST_OP_SUB,
    AST_OP_MUL,
    AST_OP_DIV,
    AST_OP_MOD,
    AST_OP_BIT_AND,
    AST_OP_BIT_OR,
    AST_OP_BIT_XOR,
    AST_OP_LOG_AND,
    AST_OP_LOG_OR,
    AST_OP_LEFT_SHIFT,
    AST_OP_RIGHT_SHIFT,
    AST_OP_CMP,
    AST_OP_COMPOSE,
};

class BinaryOp : public Node {
protected:
    OpType _op_type;
    Node* _left;
    Node* _right;
    
public:
    BinaryOp(OpType op_type, Node* left, Node* right) : 
        _op_type(op_type), _left(left), _right(right) {
        if (left) left->set_parent(this);
        if (right) right->set_parent(this);
    }

    virtual Node* accept(Visitor* v);

    void set_left(Node* left);
    void set_right(Node* right);

    OpType op_type()    { return _op_type; }
    Node* left()        { return _left; }
    Node* right()       { return _right; }
    virtual string to_string() {return "BinaryOp";}
};

class ListNode : public Node {
private:
    vector<Node*>* _node_list;

public:
    ListNode();
    
    void add(Node* n);
    
    virtual Node* accept(Visitor* v);

    vector<Node*>* node_list()  { return _node_list; }
    int size()                  { return _node_list->size(); }

    vector<Node*>::iterator begin() { return _node_list->begin(); }
    vector<Node*>::iterator end()   { return _node_list->end(); }

    Node* get(int index)        { return (*_node_list)[index]; }
    void remove(Node* n);
    void insert(Node* n, Node* t);
};

class StmtsNode : public ListNode {
public:
    StmtsNode();

    virtual Node* accept(Visitor* v);
};

class ArgsList : public ListNode {

};

class CmpNode : public BinaryOp {
private:
    int _cmp_op;

public:
    const static int CMP_LT = 0;
    const static int CMP_LE = 1;
    const static int CMP_EQU = 2;
    const static int CMP_NE = 3;
    const static int CMP_GT = 4;
    const static int CMP_GE = 5;


    CmpNode(int cmp_op) : BinaryOp(AST_OP_CMP, NULL, NULL), _cmp_op(cmp_op) {
    }

    CmpNode(Node* left, Node* right) : _cmp_op(-1), BinaryOp(AST_OP_CMP, left, right) {
    }

    void set_cmp_op(int cmp_op) { _cmp_op = cmp_op; }
    void set_left(Node* left) { _left = left; }
    void set_right(Node* right) { _right = right; }

    int cmp_op()        { return _cmp_op; }
    virtual Node* accept(Visitor* v);
};

class ConstInt : public Node {
private:
    int _value;

public:
    ConstInt(int v) : _value(v) {}

    virtual Node* accept(Visitor* v);

    virtual Type* get_type(TypeTable* type_table);
    virtual string to_string() {return "ConstInt(" + std::to_string(_value) + ")";}
    int value() { return _value; }
};

class ConstDouble : public Node {
private:
    double _value;

public:
    ConstDouble(double v) : _value(v) {}

    virtual Node* accept(Visitor* v);

    virtual Type* get_type(TypeTable* type_table);
    virtual string to_string() {return "ConstDouble(" + std::to_string(_value) + ")";}
    double value() { return _value; }
};

class ConstArray : public Node {
private:
    ListNode* _arr;
public:
    ConstArray(ListNode* v) : _arr(v) {}

    virtual Node* accept(Visitor* v);

    ListNode* arr() { return _arr; }
};

class ConstString : public Node {
public:
    char*   _value;
    int     _length;

    ConstString(int length, const char* s);

    ~ConstString();

    virtual Node* accept(Visitor* v);

    virtual Type* get_type(TypeTable* type_table);
};

class LogicNotNode: public Node {
private:
    Node* _value;
public:
    LogicNotNode(Node* v) : _value(v) {
    }

    Node* value() { return _value; }
    virtual Node* accept(Visitor* v);
};

class VarNode : public Node {
private:
    string  _name;
    ListNode*   _poly_args;

public:
    VarNode(Token* name);
    VarNode(const char* name) : _name(name) {}
    
    virtual Node* accept(Visitor* v);
    string& name()        { return _name; }
    virtual string to_string() { return "VarNode(" + _name + ")"; }

    void set_poly_args(ListNode* n) { _poly_args = n; }
    ListNode* poly_args()                { return _poly_args; }
};

class CallNode : public Node {
private:
    Node*       _name;
    ListNode*   _args;

public:
    CallNode(Node* name, ListNode* args) : _name(name), _args(args) {}
    
    virtual Node* accept(Visitor* v);

    Node* callee()      { return _name; }
    ListNode* args()    { return _args; }

    void set_callee(Node* n)   { _name = n; }
};

class SubscrNode : public Node {
private:
    Node*       _name;
    Node*       _index;

public:
    SubscrNode(Node* name, Node* i) : _name(name), _index(i) {}
    
    virtual Node* accept(Visitor* v);

    Node* name()      { return _name; }
    Node* index()     { return _index; }

    void set_name(Node* n) { _name = n; }
    void set_index(Node* i) { _index = i; }
};

class DotNode : public Node {
private:
    Node* _name;
    Node* _field;

public:
    DotNode(Node* name, Node* field) : _name(name), _field(field) {}

    virtual Node* accept(Visitor* v);

    Node* name()      { return _name; }
    Node* field()     { return _field; }
};

class SequenceNode : public Node {
private:
    Node* _left;
    Node* _right;

public:
    SequenceNode(Node* left, Node* right) : _left(left), _right(right) {}

    virtual Node* accept(Visitor* v);

    Node* left()        { return _left; }
    Node* right()       { return _right; }
};

class AssignNode : public Node {
private:
    Node*   _left;
    Node*   _right;
    
public:
    AssignNode(Node* left, Node* right):
        _left(left), _right(right){}
        
    virtual Node* accept(Visitor* v);

    Node* left()        { return _left; }
    Node* right()       { return _right; }

    void set_left(Node* left) { _left = left; }
    void set_right(Node* right) { _right = right; }
};

class PrintNode : public Node {
private:
    Node*    _value;
    
public:
    PrintNode(Node* value) : _value(value) {}
        
    virtual Node* accept(Visitor* v);
    Node* value() { return _value; }
};

class ReturnNode : public Node {
private:
    Node*    _value;

public:
    ReturnNode(Node* value) : _value(value) {}

    virtual Node* accept(Visitor* v);
    Node* value() { return _value; }
};

class IfNode : public Node {
private:
    Node*   _cond;
    Node*   _then;
    Node*   _else;

public:
    IfNode() : _cond(nullptr), _then(nullptr), _else(nullptr) {}

    IfNode(Node* cond, Node* then_part, Node* else_part) : 
        _cond(cond), _then(then_part), _else(else_part) {}

    virtual Node* accept(Visitor* v);

    Node* cond()        { return _cond; }
    Node* then_part()   { return _then; }
    Node* else_part()   { return _else; }

    void set_cond(Node* cond) { _cond = cond; }
    void set_then_part(Node* then_part) { _then = then_part; }
    void set_else_part(Node* else_part) { _else = else_part; }
};

class WhileNode : public Node {
private:
    Node*   _cond;
    Node*   _body;

public:
    WhileNode() : _cond(nullptr), _body(nullptr) {}

    WhileNode(Node* cond, Node* body) :
        _cond(cond), _body(body) {}

    virtual Node* accept(Visitor* v);

    Node* cond()        { return _cond; }
    Node* body()        { return _body; }

    void set_cond(Node* cond) { _cond = cond; }
    void set_body(Node* body) { _body = body; }
};

enum TypeEnum {
    TYPE_NODE_ARR,
    TYPE_NODE_CON,
};

class ArrowTypeNode : public Node {
private:
    Node* _src;
    Node* _dst;

public:
    ArrowTypeNode(Node* src, Node* dst) : 
        _src(src), _dst(dst) {}

    Node* src()     { return _src; }
    Node* dst()     { return _dst; }

    virtual Node* accept(Visitor* v);

    virtual Type* get_type(TypeTable* type_table);
};

class ArrayTypeNode : public Node {
private:
    Node* _elm_type;
    Node* _size;

public:
    ArrayTypeNode(Node* t, Node* s) : 
        _elm_type(t), _size(s) {}

    Node* elm_type()    { return _elm_type; }
    Node* size()        { return _size; }

    virtual Node* accept(Visitor* v);
};

class VarDefNode : public Node {
private:
    string _name;
    Node*  _def_type;
    Node*  _init_value;

public:
    VarDefNode(string name) : _name(name) {}

    VarDefNode(string name, Node* var_type, Node* init_value) :
        _name(name), _def_type(var_type), _init_value(init_value) {}

    virtual Node* accept(Visitor* v);

    void set_name(string name) { _name = name; }
    void set_def_type(Node* dt)  { _def_type = dt; }
    void set_init_value(Node* v) { _init_value = v; }

    const string& name()    { return _name; }
    Node* def_type()        { return _def_type; }
    Node* init_value()      { return _init_value; }
};

class ParamDefNode : public Node {
private:
    string      _name;
    Node*   _type;
    Node*       _default;

public:
    ParamDefNode(string name, Node* var_type, Node* init_value) :
        _name(name), _type(var_type), _default(init_value) {}

    virtual Node* accept(Visitor* v);

    const string& name()        { return _name; }
    Node* def_type()            { return _type; }
    Node* default_value()       { return _default; }
};

class FuncDefNode : public Node {
private:
    string      _name;
    Node*       _return_type_node;
    ListNode*   _args;
    Node*       _body;
    ListNode*   _poly_params;

    ListNode*   _generic_args;

    SymbolTable*    _sym_table;
    Type*           _ret_type;

public:
    FuncDefNode(Token* t);
    FuncDefNode(string& name, ListNode* args, Node* ret_tp, Node* body): 
        _name(name), _args(args), _return_type_node(ret_tp), _body(body) {
        _ret_type = nullptr;
        _generic_args = nullptr;
        _sym_table = nullptr;
        _poly_params = nullptr;
    }

    virtual Node* accept(Visitor* v);

    const string& name()            { return _name; }
    ListNode* args()                { return _args; }
    void set_args(ListNode* args)   { _args = args; }
    Node* body()                    { return _body; }
    void set_body(Node* body)       { _body = body; }
    Node* ret_type_node()           { return _return_type_node; }
    void set_ret_type_node(Node* t) { _return_type_node = t; }
    Type* ret_type()                { return _ret_type; }
    void set_ret_type(Type* t)      { _ret_type = t; }
    void set_generic_args(ListNode* args)   { _generic_args= args; }
    ListNode* generic_args()                { return _generic_args; }
    bool has_generic_args()                 { return _generic_args != nullptr && _generic_args->size() > 0; }

    SymbolTable* sym_table()            { return _sym_table; }
    void set_sym_table(SymbolTable* st) { _sym_table = st; }

    TypeTable* type_table();
};

class RecordDefNode : public Node {
private:
    string _name;
    ListNode* _fields;
    ListNode* _generic_args;
    ListNode* _methods;

    SymbolTable*    _sym_table;

public:
    RecordDefNode(string name, ListNode* fields, ListNode* generic_args) : _name(name), _fields(fields), _generic_args(generic_args) {
        _sym_table = nullptr;
        _methods = nullptr;
    }

    virtual Node* accept(Visitor* v);

    const string& name()        { return _name; }
    ListNode* fields()          { return _fields; }
    void set_fields(ListNode* fields)   { _fields = fields; }
    ListNode* generic_args()    { return _generic_args; }
    void set_generic_args(ListNode* args)   { _generic_args = args; }
    bool has_generic_args()                 { return _generic_args != nullptr && _generic_args->size() > 0; }

    SymbolTable* sym_table()            { return _sym_table; }
    void set_sym_table(SymbolTable* st) { _sym_table = st; }

    ListNode* methods() { return _methods; }
    void set_methods(ListNode* methods) { _methods = methods; }
    void add_method(FuncDefNode* method);
};

class LambdaDef : public Node {
private:
    Node*           _param;
    Node*           _body;

    SymbolTable*    _sym_table;

public:
    LambdaDef(Node* param, Node* body):
        _param(param), _body(body){}
    ~LambdaDef();

    Node* param() { return _param; }
    Node* body()  { return _body; }

    SymbolTable* sym_table()            { return _sym_table; }
    void set_sym_table(SymbolTable* st) { _sym_table = st; }

    virtual Node* accept(Visitor* visitor);
};

class TupleNode : public Node {
private:
    ListNode* _args;

public:
    TupleNode(ListNode* args) : _args(args) {}

    virtual Node* accept(Visitor* v);

    ListNode* args() { return _args; }
};

/*
 * Nodes for type defination.
 *
 * TypeVarNode: 代表一个单名字的类型，例如 Int, String, U, T 等
 */
class TypeVarNode : public Node {
private:
    string  _name;

public:
    TypeVarNode(string name, Node* type) : _name(name) {}

    virtual Node* accept(Visitor* v);
    string& name()        { return _name; }

    virtual Type* get_type(TypeTable* type_table);
};

class TypeVarDefNode : public Node {
private:
    string  _name;
    Node*   _value;

public:
    TypeVarDefNode(string& name, Node* value) : _name(name), _value(value) {}
    ~TypeVarDefNode();

    string& name()          { return _name; }
    Node* value()            { return _value; }

    virtual Node* accept(Visitor* visitor);
};

class TypeFuncDefNode : public Node {
private:
    string  _name;
    Node*   _body;
    ListNode*   _args;

public:
    TypeFuncDefNode(string& name) : _name(name) {}
    ~TypeFuncDefNode();

    string& name()              { return _name; }
    Node* body()                { return _body; }
    void set_body(Node* b)      { _body = b; }
    ListNode* args()            { return _args; }
    void set_args(ListNode* a)  { _args = a; }

    virtual Node* accept(Visitor* visitor);
};

class TypeCallNode : public Node {
private:
    Node*       _name;
    ListNode*   _args;

public:
    TypeCallNode(Node* name, ListNode* args) : _name(name), _args(args) {}

    virtual Node* accept(Visitor* v);

    Node* callee()      { return _name; }
    ListNode* args()    { return _args; }

    virtual Type* get_type(TypeTable* type_table);
};

class TypeTupleNode : public Node {
private:
    ListNode* _args;

public:
    TypeTupleNode(ListNode* args) : _args(args) {}

    virtual Node* accept(Visitor* v);

    ListNode* args()    { return _args; }

    virtual Type* get_type(TypeTable* type_table);
};

class TypeConstInt : public Node {
public:
    int value;

    TypeConstInt(int v) : value(v) {}

    virtual Node* accept(Visitor* v);
};

class BuiltinConst : public Node {
public:
    int _builtin_id;

    static const int UNIT = 0;
    static const int TRUE = 1;
    static const int FALSE = 2;
    
    static BuiltinConst* UnitNode;
    static BuiltinConst* TrueNode;
    static BuiltinConst* FalseNode;
    
    BuiltinConst(int v) : _builtin_id(v) {}
    
    virtual Node* accept(Visitor* v);
};

class ImportNode : public Node {
private:
    vector<string> _packages;
    string _name;
    string _alias;

public:
    ImportNode() {}

    virtual Node* accept(Visitor* v);

    void add_package(string package) { _packages.push_back(package); }
    void set_name(string name) { _name = name; }
    void set_alias(string alias) { _alias = alias; }
    void set_packages(vector<string> packages) { _packages = packages; }

    vector<string>& packages() { return _packages; }
    string& name() { return _name; }
    string& alias() { return _alias; }
};

class ImplementNode : public Node {
private:
    string _trait_name;
    string _for_name;
    ListNode* _methods;
    ListNode* _generic_args;

public:
    ImplementNode(string trait_name, string for_name, ListNode* methods, ListNode* generic_args) :
        _trait_name(trait_name), _for_name(for_name), _methods(methods), _generic_args(generic_args) {}

    virtual Node* accept(Visitor* v);

    string trait_name() { return _trait_name; }
    string for_name() { return _for_name; }
    ListNode* methods() { return _methods; }
    ListNode* generic_args() { return _generic_args; }
    void set_generic_args(ListNode* args) { _generic_args = args; }
};

class TraitDefNode : public Node {
private:
    string _name;
    StmtsNode* _methods;
    SymbolTable* _sym_table;

    ListNode* _generic_args;

public:
    TraitDefNode(string name, StmtsNode* methods) : _name(name), _methods(methods) {}

    virtual Node* accept(Visitor* v);

    string name() { return _name; }
    StmtsNode* methods() { return _methods; }

    SymbolTable* sym_table()            { return _sym_table; }
    void set_sym_table(SymbolTable* st) { _sym_table = st; }

    ListNode* generic_args() { return _generic_args; }
    void set_generic_args(ListNode* args) { _generic_args = args; }
    bool has_generic_args() { return _generic_args != nullptr && _generic_args->size() > 0; }
};

#endif
